"""export"""
import os
import sys
import argparse
import numpy as np
from mindspore import Tensor, export, load_checkpoint, context
import mindspore as ms
from mind3d.utils.load_yaml import load_yaml
from mind3d.utils.synergynet_builder import build_model

def export(opt):
    model = build_model(img_size=opt['export']['img_size'], mode="test")
    if not os.path.exists('./mindir'):
        os.mkdir('./mindir')
    load_checkpoint(opt['export']['model'], net=model)
    input_data = Tensor(np.random.uniform(-1.0, 1.0, size=[1, 3, 120, 120]).astype(np.float32), ms.float32)
    export(model, input_data, file_name='./mindir/synergynet', file_format=opt['export']['file_format'])
    print("successfully export model")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='SynergyNet model export')
    parser.add_argument('-opt', type=str, default='mind3d/configs/synergy_net/synergynet.yaml',
                            help='Path to option YAML file.')
    args = parser.parse_args()[0]
    opt = load_yaml(args.opt)
    context.set_context(max_call_depth=100000, mode=context.GRAPH_MODE, device_target=opt['export']['device_target'], device_id=opt['export']['device_id'])
            
    export(opt)
